/**
 * 
 */
package edu.umd.clip.lm.model.data;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.channels.FileChannel;

import edu.umd.clip.lm.questions.Question;
import edu.umd.clip.lm.util.Long2IntMap;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class TrainingDataUtil {

	public static ReadableTrainingData makeContextReducedTrainingData(ReadableTrainingData input, 
			int overtOrder, int hiddenOrder) 
	{
		return makeContextReducedTrainingData(input, overtOrder, hiddenOrder, null);
	}
	
	public static ReadableTrainingData makeContextReducedTrainingData(ReadableTrainingData input, 
			int overtOrder, int hiddenOrder, final TrainingDataFilter filter) 
	{
		TrainingDataReader proxyTrainReader = new ProxyTrainingDataReader(input);

		if (filter != null) {
			input = new FilteredReadableTrainingData(proxyTrainReader, filter);
			proxyTrainReader = new ProxyTrainingDataReader(input);
		}
		
		input = new ContextReducingReadableTrainingData(proxyTrainReader, overtOrder, hiddenOrder, 
					TrainingDataBlock.getMaxCapacity() * 2);

		return input;
	}
	
	public static void reduceContext(ReadableTrainingData input, WritableTrainingData output, 
			int overtOrder, int hiddenOrder, 
			final TrainingDataFilter filter, final Long2IntMap counts) throws IOException 
	{
		input = makeContextReducedTrainingData(input, overtOrder, hiddenOrder, filter);
		
		input.start();
		output.setContextSize(input.getContextSize());
		output.start();
		for(; input.hasNext();) {
			TrainingDataBlock block = input.next();
			if (counts != null) {
				block.addCounts(counts);
			}
			output.add(block);
			//System.err.printf("ContextReduce: writing block, totalCount = %d\n", block.getTotalCount());
		}
	}

	public static void reduceContext(ReadableTrainingData input, WritableTrainingData output, 
			int overtOrder, int hiddenOrder, 
			TrainingDataFilter filter) throws IOException 
	{
		reduceContext(input, output, overtOrder, hiddenOrder, filter, null);
	}
	
	public static void reduceContext(ReadableTrainingData input, WritableTrainingData output, 
			int overtOrder, int hiddenOrder, 
			final Long2IntMap counts) throws IOException 
	{
		reduceContext(input, output, overtOrder, hiddenOrder, null, counts);
	}
	
	public static void reduceContext(ReadableTrainingData input, WritableTrainingData output, 
			int overtOrder, int hiddenOrder) throws IOException 
	{
		reduceContext(input, output, overtOrder, hiddenOrder, null, null);
	}

	public static void combineAndReduceContext(String[] inputFiles, String outputFile, 
			int overtOrder, int hiddenOrder, 
			int clusterid, TrainingDataFilter filter) 
	throws IOException 
	{
		RandomAccessFile outFile = new RandomAccessFile(new File(outputFile), "rw");
		outFile.getChannel().truncate(0);
		
		OnDiskTrainingDataNodeWriter writer = new OnDiskTrainingDataNodeWriter(outFile.getChannel());
		
		WritableTrainingDataNode node = writer.createNode(clusterid, 1);
		writer.add(node);
		
		// Aggregate multiple files into one
		for(String fname : inputFiles) {
			FileChannel channel = new FileInputStream(fname).getChannel();
			TrainingDataReader reader = new OnDiskTrainingDataReader(channel);
			ReadableTrainingData smoothingData = new ReadableTrainingData(reader);

			TrainingDataUtil.reduceContext(smoothingData, node.getData(0), overtOrder, hiddenOrder, filter);
			
			channel.close();
		}
		node.getData(0).finish();
		writer.close();
		
		outFile.close();
	}
	
	public static void splitData(ReadableTrainingData data, Question question, 
			WritableTrainingData trueData, WritableTrainingData falseData) throws IOException 
	{
		splitData(data, question, trueData, falseData, null);
	}
	
	public static void splitData(ReadableTrainingData data, Question question, 
			WritableTrainingData trueData, WritableTrainingData falseData,
			final Long2IntMap counts) throws IOException 
	{
		splitData(data, question, trueData, falseData, counts, null, null);
	}	
	
	public static void splitData(ReadableTrainingData data, Question question, 
			WritableTrainingData trueData, WritableTrainingData falseData,
			final Long2IntMap counts, final Long2IntMap trueCounts, final Long2IntMap falseCounts) 
	throws IOException 
	{
		data.start();
		trueData.setContextSize(data.getContextSize());
		falseData.setContextSize(data.getContextSize());
		
		trueData.start();
		falseData.start();
		
		TrainingDataBlock trueBlock = new TrainingDataBlock();
		TrainingDataBlock falseBlock = new TrainingDataBlock();
		
		for(; data.hasNext(); ) {
			TrainingDataBlock block = data.next();
			if (block.hasData()) {
				if (counts != null) {
					block.addCounts(counts);
				}

				for(ContextFuturesPair pair : block) {
					if (question.test(pair.getContext())) {
						if (!trueBlock.add(pair)) {
							trueData.add(trueBlock);
							if (trueCounts != null) {
								trueBlock.addCounts(trueCounts);
							}
							trueBlock = new TrainingDataBlock();
							trueBlock.add(pair);
						}
					} else {
						if (!falseBlock.add(pair)) {
							falseData.add(falseBlock);
							if (falseCounts != null) {
								falseBlock.addCounts(falseCounts);
							}
							falseBlock = new TrainingDataBlock();
							falseBlock.add(pair);
						}
					}
				}
			}
		}
		if (trueBlock.hasData()) {
			trueData.add(trueBlock);
			if (trueCounts != null) {
				trueBlock.addCounts(trueCounts);
			}
		}
		if (falseBlock.hasData()) {
			falseData.add(falseBlock);
			if (falseCounts != null) {
				falseBlock.addCounts(falseCounts);
			}
		}
		trueData.finish();
		falseData.finish();
	}

}

